import random
from torch.autograd import Variable


class Agent(object):
    def __init__(self, model, action_size):
        self.model = model
        self.action_size = action_size

    def _select_action(self, state, eps_threshold=0.05):
        sample = random.random()

        if sample > eps_threshold:
            actions_value = self.model.forward(Variable(state)).data
            return actions_value.max(1)[1][0]
        else:
            if random.random() >= 0.3:
                return 1
            return 0

    def select_action(self, state, eps_threshold=0.05):
        action = self._select_action(state, eps_threshold)
        return action
